import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 假设我们有一些二维数据点 (x, y, z)
# 在这里，我们生成一些带有噪声的模拟数据

x = np.linspace(-512, 511, 1024)
y = np.linspace(-640, 639, 1280)
X, Y = np.meshgrid(x, y)

att_4d = np.load('att_4d.npy', allow_pickle=True)
Z = att_4d[1, -1]
# 将数据点展平为一维数组，以便用于curve_fit
x_flat = X.flatten()
y_flat = Y.flatten()
z_flat = Z.flatten()


# 定义抛物面函数
def paraboloid(variables, a, b, c, d, e, f):
    x, y = variables
    return a * x ** 2 + b * y ** 2 + c * x * y + d * x + e * y + f


# 使用curve_fit进行拟合
popt, pcov = curve_fit(paraboloid, (x_flat, y_flat), z_flat, p0=[1, 1, 0, 0, 0, 0])

# 提取拟合参数
a, b, c, d, e, f = popt
print('a = ', a)
print('b = ', b)
print('c = ', c)
print('d = ', d)
print('e = ', e)
print('f = ', f)

# 创建用于绘图的网格
X_fit = np.linspace(x.min(), x.max(), 100)
Y_fit = np.linspace(y.min(), y.max(), 100)
X_fit, Y_fit = np.meshgrid(X_fit, Y_fit)
Z_fit = paraboloid((X_fit, Y_fit), *popt)

# 绘制原始数据点和拟合的抛物面
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
# ax.scatter(x_flat, y_flat, z_flat, c='red', label='Data points')
ax.plot_surface(X_fit, Y_fit, Z_fit, cmap='viridis', label='Fitted paraboloid')

# 添加坐标轴标签和图例
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.legend()

# 显示图形
plt.show()